import pandas as pd
import chromadb
from chromadb.config import Settings
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings

import os
from dotenv import load_dotenv
load_dotenv()
API_KEY = os.getenv("OPENAI_API_KEY")
BASE_URL = os.getenv("OPENAI_PROXY")

# 读取 Excel 文件
excel_file = pd.ExcelFile('../../asset/Cars.xlsx')

df = excel_file.parse('Sheet1')


# 假设所有列都作为文本内容，将数据转为适合的文本列表格式
def preprocess_car_data(df):
    """优化汽车数据预处理：读取表格数据并结构化文本"""
    texts = []
    for index, row in df.iterrows():
        # 提取关键字段（根据实际数据字段名映射）
        brand = row.get('brand_name', '')  # 汽车品牌
        model = row.get('car_name', '')  # 车型配置（原代码中model未定义，此处补充）
        series = row.get('series_name', '')  # 车系名称
        price = row.get('sh_price', '')  # 实际售价（使用sh_price作为价格）
        year = row.get('car_year', '')  # 车辆年份
        car_type = row.get('car_type', '')  # 车型类型
        fuel_type = row.get('fuel_form', '')  # 燃料类型
        driver_form = row.get('driver_form', '')  # 驱动方式

        # 重新组织文本，突出关键信息
        structured_text = f"品牌:{brand} 车系:{series} 车型:{model} 价格:{price} 年份:{year} "
        structured_text += f"车型类型:{car_type} 燃料类型:{fuel_type} 驱动方式:{driver_form} "

        # 拼接所有非空字段值作为补充信息
        additional_info = " ".join([str(value) for value in row.values if pd.notna(value)])
        combined_text = structured_text + additional_info

        texts.append(combined_text)
    return texts


texts = preprocess_car_data(df)
for index, row in df.iterrows():
    combined_text = " ".join([str(value) for value in row.values if pd.notna(value)])
    texts.append(combined_text)


# 设置 Chroma 客户端
client = chromadb.PersistentClient(path="../../asset/chroma_db")


embedding_function = OpenAIEmbeddings(
    model="text-embedding-3-large",
    openai_api_key=API_KEY,
    openai_api_base=BASE_URL,  # 如果使用代理
    request_timeout=60,  # 设置超时时间
    max_retries=3  # 设置重试次数
)

# 将文本数据和对应的向量存入 Chroma 数据库
try:
    vector_store = Chroma.from_texts(
        texts,
        embedding_function,
        client=client
    )
except Exception as e:
    print(f"向量存储创建失败: {e}")
